import parser
import os
from .base import Base

class Pb2ONNX(Base):
    def __init__(self, success_list, failed_list):
        print('Pb2ONNX Init--------------------------')
        self.tool_path = os.path.join(parser.get_value('maca_converter_root'), parser.get_value('maca_converter_path'))
        self.model_path = os.path.join(parser.get_value('model_root'), parser.get_value('pb_model_path'))
        self.inputs = parser.get_value('pb_model_inputs')
        self.outputs = parser.get_value('pb_model_outputs')
        self.inputs_as_nchw = parser.get_value('pb_inputs_as_nchw')
        self.success_list = success_list
        self.failed_list = failed_list
        self.input = os.path.join(parser.get_value('model_root'), parser.get_value('pb_input_npy'))#parser.get_value('h5_input_npy')
        self.output = os.path.join(parser.get_value('model_root'), parser.get_value('pb_output_npy'))#parser.get_value('h5_output_npy')

    def test(self):
        cwd = os.getcwd()
        output = os.path.join(cwd, 'tmp.onnx')

        cmd = 'python ' + 'model_convert.py' + ' --model_path ' + self.model_path + ' --model_type tf-graph ' + '--output ' + output + ' --inputs ' + self.inputs + ' --outputs ' + self.outputs
        if self.inputs_as_nchw != '':
            cmd = cmd + ' --inputs_as_nchw ' + self.inputs_as_nchw

        print('test_pb2onnx: ', cmd)

        self.test_model2onnx(self.tool_path, 'pb2onnx', cmd, output, self.success_list, self.failed_list, self.input, self.output)